"""
固件管理模块
负责固件文件的读取、解析和管理
"""
import os
import struct
from typing import Tuple, List, Optional, Dict
from dataclasses import dataclass


@dataclass
class FirmwareSegment:
    """固件段数据结构"""
    address: int           # 起始地址
    size: int              # 数据大小
    data: list[int]        # 数据内容（字节列表）
    

class FirmwareManager:
    """固件管理器"""
    
    def __init__(self):
        self._segments: List[FirmwareSegment] = []
        self._loaded_file = ""
        self._total_size = 0
        
    def load_srec_file(self, file_path: str) -> Tuple[bool, str]:
        """
        加载S-Record格式的固件文件
        
        Args:
            file_path: 固件文件路径
            
        Returns:
            Tuple[bool, str]: (是否成功, 消息)
        """
        try:
            if not os.path.exists(file_path):
                return False, f"固件文件不存在: {file_path}"
            
            # 清空之前的数据
            self._segments.clear()
            self._total_size = 0
            
            # 临时存储地址连续的数据
            temp_data: Dict[int, int] = {}  # address -> data_byte
            
            with open(file_path, 'r') as f:
                line_num = 0
                for line in f:
                    line_num += 1
                    line = line.strip()
                    
                    # 跳过空行
                    if not line:
                        continue
                    
                    # 检查S-Record格式
                    if not line.startswith('S'):
                        continue
                    
                    # 解析S-Record
                    suc, msg = self._parse_srec_line(line, temp_data)
                    if not suc:
                        return False, f"解析第{line_num}行失败: {msg}"
            
            # 将连续的数据组合成段
            self._group_data_to_segments(temp_data)
            
            self._loaded_file = file_path
            return True, f"成功加载固件文件: {file_path}, 共{len(self._segments)}个段, 总大小: {self._total_size} bytes"
            
        except Exception as e:
            return False, f"加载固件文件异常: {str(e)}"
    
    def _parse_srec_line(self, line: str, temp_data: Dict[int, int]) -> Tuple[bool, str]:
        """解析S-Record行"""
        try:
            if len(line) < 4:
                return False, "行长度不足"
            
            record_type = line[1]
            
            # 只处理S1, S2, S3数据记录
            if record_type not in ['1', '2', '3']:
                return True, "跳过非数据记录"
            
            # 获取数据长度
            byte_count = int(line[2:4], 16)
            
            # 计算地址长度
            addr_len = 2 if record_type == '1' else (3 if record_type == '2' else 4)
            
            # 检查行长度
            expected_len = 4 + 2 * byte_count
            if len(line) != expected_len:
                return False, f"行长度错误: {len(line)} != {expected_len}"
            
            # 解析地址
            addr_hex = line[4:4+2*addr_len]
            address = int(addr_hex, 16)
            
            # 解析数据
            data_start = 4 + 2 * addr_len
            data_end = data_start + 2 * (byte_count - addr_len - 1)  # 减去地址长度和校验和
            
            for i in range(data_start, data_end, 2):
                data_byte = int(line[i:i+2], 16)
                temp_data[address + (i - data_start) // 2] = data_byte
            
            return True, "解析成功"
            
        except Exception as e:
            return False, f"解析异常: {str(e)}"
    
    def _group_data_to_segments(self, temp_data: Dict[int, int]) -> None:
        """将连续的数据组合成段"""
        if not temp_data:
            return
        
        # 按地址排序
        sorted_addresses = sorted(temp_data.keys())
        
        # 组合连续的数据
        current_start = sorted_addresses[0]
        current_data = []
        
        for addr in sorted_addresses:
            if current_data and addr != current_start + len(current_data):
                # 地址不连续，创建新段
                segment = FirmwareSegment(
                    address=current_start,
                    size=len(current_data),
                    data=current_data.copy()
                )
                self._segments.append(segment)
                self._total_size += len(current_data)
                
                # 开始新段
                current_start = addr
                current_data = []
            
            current_data.append(temp_data[addr])
        
        # 处理最后一段
        if current_data:
            segment = FirmwareSegment(
                address=current_start,
                size=len(current_data),
                data=current_data.copy()
            )
            self._segments.append(segment)
            self._total_size += len(current_data)
    
    def load_hex_file(self, file_path: str) -> Tuple[bool, str]:
        """
        加载Intel HEX格式的固件文件
        
        Args:
            file_path: 固件文件路径
            
        Returns:
            Tuple[bool, str]: (是否成功, 消息)
        """
        try:
            if not os.path.exists(file_path):
                return False, f"固件文件不存在: {file_path}"
            
            # 清空之前的数据
            self._segments.clear()
            self._total_size = 0
            
            # 临时存储地址连续的数据
            temp_data: Dict[int, int] = {}  # address -> data_byte
            base_address = 0
            
            with open(file_path, 'r') as f:
                line_num = 0
                for line in f:
                    line_num += 1
                    line = line.strip()
                    
                    # 跳过空行
                    if not line:
                        continue
                    
                    # 检查Intel HEX格式
                    if not line.startswith(':'):
                        continue
                    
                    # 解析Intel HEX
                    suc, base_addr, msg = self._parse_hex_line(line, temp_data, base_address)
                    if not suc:
                        return False, f"解析第{line_num}行失败: {msg}"
                    
                    if base_addr is not None:
                        base_address = base_addr
            
            # 将连续的数据组合成段
            self._group_data_to_segments(temp_data)
            
            self._loaded_file = file_path
            return True, f"成功加载固件文件: {file_path}, 共{len(self._segments)}个段, 总大小: {self._total_size} bytes"
            
        except Exception as e:
            return False, f"加载固件文件异常: {str(e)}"
    
    def _parse_hex_line(self, line: str, temp_data: Dict[int, int], base_address: int) -> Tuple[bool, Optional[int], str]:
        """解析Intel HEX行"""
        try:
            if len(line) < 11:
                return False, None, "行长度不足"
            
            # 解析基本字段
            byte_count = int(line[1:3], 16)
            address = int(line[3:7], 16)
            record_type = int(line[7:9], 16)
            
            # 计算期望长度
            expected_len = 11 + 2 * byte_count
            if len(line) != expected_len:
                return False, None, f"行长度错误: {len(line)} != {expected_len}"
            
            # 处理不同类型的记录
            if record_type == 0x00:  # 数据记录
                # 解析数据
                for i in range(byte_count):
                    data_byte = int(line[9+2*i:11+2*i], 16)
                    full_address = base_address + address + i
                    temp_data[full_address] = data_byte
                    
                return True, None, "数据记录解析成功"
                
            elif record_type == 0x01:  # 文件结束记录
                return True, None, "文件结束记录"
                
            elif record_type == 0x04:  # 扩展线性地址记录
                extended_addr = int(line[9:13], 16)
                new_base_address = extended_addr << 16
                return True, new_base_address, "扩展地址记录解析成功"
                
            else:
                return True, None, f"跳过记录类型: 0x{record_type:02X}"
                
        except Exception as e:
            return False, None, f"解析异常: {str(e)}"
    
    def get_segment_count(self) -> int:
        """获取固件段数量"""
        return len(self._segments)
    
    def get_segment(self, index: int) -> Tuple[bool, Optional[FirmwareSegment], str]:
        """
        获取指定索引的固件段
        
        Args:
            index: 段索引
            
        Returns:
            Tuple[bool, Optional[FirmwareSegment], str]: (是否成功, 段数据, 消息)
        """
        if index < 0 or index >= len(self._segments):
            return False, None, f"段索引越界: {index}"
        
        return True, self._segments[index], "获取段成功"
    
    def get_segments(self) -> List[FirmwareSegment]:
        """获取所有固件段"""
        return self._segments.copy()
    
    def get_total_size(self) -> int:
        """获取固件总大小"""
        return self._total_size
    
    def get_loaded_file(self) -> str:
        """获取已加载的文件路径"""
        return self._loaded_file
    
    def is_loaded(self) -> bool:
        """检查是否已加载固件"""
        return len(self._segments) > 0
    
    def clear(self) -> None:
        """清空固件数据"""
        self._segments.clear()
        self._loaded_file = ""
        self._total_size = 0
    
    def get_address_range(self) -> Tuple[int, int]:
        """
        获取固件地址范围
        
        Returns:
            Tuple[int, int]: (最小地址, 最大地址)
        """
        if not self._segments:
            return 0, 0
        
        min_addr = min(seg.address for seg in self._segments)
        max_addr = max(seg.address + seg.size - 1 for seg in self._segments)
        
        return min_addr, max_addr
    
    def get_data_at_address(self, address: int, size: int) -> Tuple[bool, Optional[list[int]], str]:
        """
        获取指定地址的数据
        
        Args:
            address: 起始地址
            size: 数据大小
            
        Returns:
            Tuple[bool, Optional[list[int]], str]: (是否成功, 数据, 消息)
        """
        try:
            result_data = []
            
            for i in range(size):
                current_addr = address + i
                found = False
                
                # 在所有段中查找地址
                for segment in self._segments:
                    if segment.address <= current_addr < segment.address + segment.size:
                        offset = current_addr - segment.address
                        result_data.append(segment.data[offset])
                        found = True
                        break
                
                if not found:
                    return False, None, f"地址 0x{current_addr:08X} 不在固件范围内"
            
            return True, result_data, "获取数据成功"
            
        except Exception as e:
            return False, None, f"获取数据异常: {str(e)}"
    
    def get_all_firmware_data(self) -> list[int]:
        """
        获取所有固件段的数据（按地址排序）
        
        Returns:
            list[int]: 所有固件数据的列表
        """
        all_data = []
        
        # 按地址排序段
        sorted_segments = sorted(self._segments, key=lambda seg: seg.address)
        
        for segment in sorted_segments:
            all_data.extend(segment.data)
            
        return all_data


def firm_checksum(data: list[int]) -> int:
    """
    计算固件数据的校验和（上位机计算）
    
    使用简单的累加校验算法，与MCU端保持一致
    
    Args:
        data: 所有代码数据的列表
        
    Returns:
        int: 校验和值
    """
    checksum = 0
    print(len(data))
    # 简单累加校验（32位）
    for byte_val in data:
        checksum = (checksum + byte_val) 
    
    checksum = checksum & 0xFF
    return checksum 